package com.leu;

import com.trendrr.nsq.*;
import com.trendrr.nsq.lookup.NSQLookupDynMapImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.UnsupportedEncodingException;
import java.util.Date;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

public class SpeedTest {

    protected static Logger log = LoggerFactory.getLogger(SpeedTest.class);
    static String address;
    static int port;

    /**
     * @param args
     * @throws UnsupportedEncodingException
     */
    public static void main(String[] args) throws UnsupportedEncodingException {

        address = "192.168.1.50";
        port = 4150;
        if (args.length > 0) {
            address = args[0];
            if (args.length > 1) {
                port = Integer.parseInt(args[1]);
            }
        }
        //testProduceBatch();
       testConsume();
    }

    public static void testConsume() {
        final int iterations = 100000;
        final Date start = new Date();
        final AtomicInteger processed = new AtomicInteger(0);

        NSQLookup lookup = new NSQLookupDynMapImpl();
        lookup.addAddr("192.168.1.50", 4161);
        NSQConsumer consumer = new NSQConsumer(lookup, "speedtest", "testconsumer", new NSQMessageCallback() {

            public void message(NSQMessage message) {
                try {
                    message.finished();
                    int p = processed.incrementAndGet();
                    String mess = new String(message.getMessage(), "UTF-8");

                    if (p % 1000 == 0) {
                        System.out.println(p + "接收消息consumer: " + mess);
                    }
                    if (p == iterations) {
                        System.out.println("completed " + iterations + " in " + (new Date().getTime() - start.getTime()));
                    }

                } catch (Exception e) {
                    log.error("Caught", e);
                }
            }
            public void error(Exception x) {
                log.warn("Caught", x);
            }
        });

        System.out.println("Starting consumer");
        consumer.start();

        //TODO: have a wait latch

    }
    public static void testProduceBatch() {
        int iterations = 100000;  //迭代次数
        int batchsize = 100;  //批量大小
        int poolsize = 1;  //池大小
        int messageSize = 10000; //消息字节数.
        NSQProducer producer = new NSQProducer().addAddress(address, port, poolsize);
        producer.configureBatch("speedtest",
                new BatchCallback() {
                    public void batchSuccess(String topic, int num) {
                    }

                    public void batchError(Exception ex, String topic, List<byte[]> messages) {
                        ex.printStackTrace();
                    }
                },
                batchsize,
                null, //use default maxbytes
                null //use default max seconds
        );
        producer.start();
        Date start = new Date();
        StringBuilder str = new StringBuilder();
        for (int i = 0; i < messageSize; i++) {
            str.append("a");
        }
        byte[] msg = str.toString().getBytes();
        for (int i = 0; i < iterations; i++) {
            if (i % 1000 == 0) {
                System.out.println("发送消息produced: " + i);
            }
            producer.produceBatch("speedtest", msg);
        }
        System.out.println("Produced " + iterations + " batch messages in " + (new Date().getTime() - start.getTime()) + " millis");
        producer.close();
    }
}